package com.xyz.module;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import scala.Tuple2;

import java.util.*;

public class TopN {
    public static void main(String []args){
        SparkConf sparkConf = new SparkConf().setAppName("Spark Application (java)").setMaster("local");

        JavaSparkContext javaSparkContext = new JavaSparkContext(sparkConf);

        String inputPath = "/Users/maxiangyang/sparktest/testtopn.txt";
        int N = 10;

        JavaRDD<String> textFile = javaSparkContext.textFile(inputPath);

        JavaPairRDD<String,Integer> javaPairRDD = textFile.flatMap(s -> Arrays.asList(s.split("\n")).iterator())
                .mapToPair(s -> {
                    if("".equals(s)){
                        s = "0";
                    }
                    return new Tuple2<>(s,Integer.parseInt(s));
                });

        JavaRDD<SortedMap<Integer,String>> singleTopN = javaPairRDD.mapPartitions(tuple2Iterator -> {
            SortedMap<Integer,String> topn = new TreeMap<>();
            while (tuple2Iterator.hasNext()){
                Tuple2<String,Integer> stringTuple2 = tuple2Iterator.next();
                topn.put(stringTuple2._2,stringTuple2._1);
                if(topn.size()>N){
                    topn.remove(topn.firstKey());
                }
            }

            return Collections.singletonList(topn).iterator();
        });

        List<SortedMap<Integer,String>> singleResult = singleTopN.collect();

        SortedMap<Integer,String> finaltopn = new TreeMap<>();
        singleResult.stream().forEach(integerStringSortedMap -> {
            integerStringSortedMap.entrySet().forEach(integerStringEntry -> {
                finaltopn.put(integerStringEntry.getKey(),integerStringEntry.getValue());
            });
        });

        finaltopn.entrySet().forEach(integerStringEntry -> {
            System.out.println(integerStringEntry.getKey()+":"+integerStringEntry.getValue());
        });


    }
}
